import torch
from torchvision.transforms import Compose, CenterCrop
from datasets.thu_chl.transforms import Repeat, toOneHot, ToTensor
from datasets.thu_chl.process import *
from datasets.thu_chl.dataset import THU_EACT_50_CHL

repr_map = {'eventFrame': get_eventFrame,
            'eventAccuFrame': get_eventAccuFrame,
            'timeSurface': get_timeSurface,
            'eventCount': get_eventCount}


def defualt_create_datasets(datafile="../THU-EACT-50-CHL", train=False, event_augmentation=False,
                            repr=['timeSurface', 'eventFrame'], time_num=9, ret_file_name=False,
                            frame_transform=None, label_transform=None, voxel=True
                            ):
    """

    Args:
        voxel: 是否采用体素编码方式
        datafile: THU-EACT-50-CHL 根路径
        train: 是否处于训练模式
        augmentation:是否对原始事件数据进行数据增强
        repr:提取事件数据的方法，包括
        repr_map = {'eventFrame':get_eventFrame,
            'eventAccuFrame':get_eventAccuFrame,
            'timeSurface':get_timeSurface,
            'eventCount':get_eventCount}
        time_num:需要提起的事件帧的数量，即为T
        ret_file_name:是否返回事件表征方式的名字
        frame_transform:对得到的事件帧进行数据处理，torchvision.transforms
        label_transform:对得到的标签进行数据处理,torchvision.transforms

    Returns:
        返回数据集格式，其中__getitem__方法返回的shape [T,len(repr),X,Y],X:[0,346], Y:[0,260]

    """
    if frame_transform is None:
        def default_frame_transform():
            return Compose([
                ToTensor(),
                # CenterCrop(256)
            ])

        frame_transform = default_frame_transform()

    if label_transform is None:
        def default_label_transform():
            return Compose([
                Repeat(time_num), toOneHot(50)
            ])

        # label_transform = default_label_transform()
        label_transform = torch.tensor

    datasets = THU_EACT_50_CHL(
        datafile,
        train=train,
        augmentation=event_augmentation,
        repr=repr,
        time_num=time_num,
        ret_file_name=ret_file_name,
        frame_transform=frame_transform,
        label_transform=label_transform,
        voxel=voxel
    )

    return datasets
